// SPDX-License-Identifier: GPL-2.0
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/kprobes.h>
#include <linux/ktime.h>
#include <linux/limits.h>
#include <linux/sched.h>
#include <linux/debugfs.h>
#include <linux/spinlock.h>
#include <linux/list.h>
#include <linux/slab.h>

#include "mod_detector.h"
#include "trap_io.h"
#include "error-inject.h"

static LIST_HEAD(emu_io_list);
static DEFINE_RAW_SPINLOCK(ei_lock);
struct ei_entry {
	struct list_head list;
	void __iomem *virt_addr;
	phys_addr_t phys_addr;
    size_t size;
};

void __iomem *find_within_emu_io(phys_addr_t addr)
{
    struct ei_entry *ent;
    void __iomem *vaddr = NULL;
    unsigned long flags;

    raw_spin_lock_irqsave(&ei_lock, flags);
    list_for_each_entry(ent, &emu_io_list, list) {
        if (addr >= ent->phys_addr && addr < ent->phys_addr + ent->size) {
            vaddr = ent->virt_addr + (addr - ent->phys_addr);
            break;
        }
    }
    raw_spin_unlock_irqrestore(&ei_lock, flags);

    return vaddr;
}

static int register_emu_io(void __iomem *virt_addr, phys_addr_t phys_addr, size_t size)
{
	struct ei_entry *ent;
    unsigned long flags;

    raw_spin_lock_irqsave(&ei_lock, flags);

	ent = kmalloc(sizeof(*ent), GFP_ATOMIC);
	if (!ent) {
		return -ENOMEM;
    }
    ent->virt_addr = virt_addr;
    ent->phys_addr = phys_addr;
    ent->size = size;
    INIT_LIST_HEAD(&ent->list);
    list_add_tail(&ent->list, &emu_io_list);

    raw_spin_unlock_irqrestore(&ei_lock, flags);

    return 0;
}

static void unregister_emu_io(void __iomem *virt_addr)
{
    struct ei_entry *ent, *n;
    unsigned long flags;

    raw_spin_lock_irqsave(&ei_lock, flags);
    list_for_each_entry_safe(ent, n, &emu_io_list, list) {
        if (ent->virt_addr == virt_addr) {
            list_del_init(&ent->list);
            kfree(ent);
        }
    }
    raw_spin_unlock_irqrestore(&ei_lock, flags);
}

static unsigned long seq_flags;
static void *ei_seq_start(struct seq_file *m, loff_t *pos)
{
    raw_spin_lock_irqsave(&ei_lock, seq_flags);
	return seq_list_start(&emu_io_list, *pos);
}

static void ei_seq_stop(struct seq_file *m, void *v)
{
    raw_spin_unlock_irqrestore(&ei_lock, seq_flags);
}

static void *ei_seq_next(struct seq_file *m, void *v, loff_t *pos)
{
	return seq_list_next(v, &emu_io_list, pos);
}

static int ei_seq_show(struct seq_file *m, void *v)
{
	struct ei_entry *ent = list_entry(v, struct ei_entry, list);

	seq_printf(m, "[address]0x%llx\t[size]0x%lx\n", ent->phys_addr, ent->size);
	return 0;
}

static const struct seq_operations ei_seq_ops = {
	.start = ei_seq_start,
	.next  = ei_seq_next,
	.stop  = ei_seq_stop,
	.show  = ei_seq_show,
};

static int ei_open(struct inode *inode, struct file *filp)
{
	return seq_open(filp, &ei_seq_ops);
}

static const struct file_operations debugfs_ei_ops = {
	.open           = ei_open,
	.read           = seq_read,
	.llseek         = seq_lseek,
	.release        = seq_release,
};

static int ei_debugfs_init(void)
{
    struct dentry *dir, *file;

    dir = debugfs_create_dir("io_emulation", NULL);
    if (!dir)
        return -ENOMEM;

    file = debugfs_create_file("list", 0444, dir, NULL, &debugfs_ei_ops);
     if (!file) {
		debugfs_remove(dir);
		return -ENOMEM;
	}

	return 0;
}

struct iomap_info {
    struct module_info *cnmod;

    phys_addr_t phys_addr;
    size_t size;
    void __iomem *virt_addr;
};

static emulation_handle_t *g_emulation_handle = NULL;

static inline void __iomem *virt2emu(void __iomem * virt)
{
    uint64_t p = (uint64_t)virt;

    return (void __iomem *)((p & 0x0000ffffffffffffL)|0xcabc000000000000L);
}

static inline void __iomem *emu2virt(void __iomem * emu)
{
    uint64_t p = (uint64_t)emu;

    return (void __iomem *)(p | 0xffff000000000000L);
}

static int ioremap_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;
    unsigned long caller_address = procedure_link_pointer(regs);

    if ((info->cnmod = lookup_module_by_text_address(caller_address)) == NULL) {
        /* The caller is not our modules, ignore */
        /* return 1, kret handler won't be called */
        return 1;
    }

    /* void __iomem *__ioremap(phys_addr_t phys_addr, size_t size, pgprot_t prot) */
    info->phys_addr = (phys_addr_t)regs_get_kernel_argument(regs, 0);
    info->size = (size_t)regs_get_kernel_argument(regs, 1);

    return 0;
}
NOKPROBE_SYMBOL(ioremap_entry);

static int ioremap_return(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;

    info->virt_addr = (void __iomem *)regs_return_value(regs);

    register_emu_io(info->virt_addr, info->phys_addr, info->size);

    /* override with return emulation address */
    regs_set_return_value(regs, (unsigned long)virt2emu(info->virt_addr));
    instruction_pointer_set(regs, procedure_link_pointer(regs));

    return 0;
}
NOKPROBE_SYMBOL(ioremap_return);

static struct kretprobe ioremap_kretprobe = {
    .kp.symbol_name = "__ioremap",
	.entry_handler	= ioremap_entry,
    .handler		= ioremap_return,
	.data_size		= sizeof(struct iomap_info),
	/* Probe up to 10 instances concurrently. */
	.maxactive		= 10,
};

static int devm_ioremap_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;
    unsigned long caller_address = procedure_link_pointer(regs);

    if ((info->cnmod = lookup_module_by_text_address(caller_address)) == NULL) {
        /* The caller is not our modules, ignore */
        /* return 1, kret handler won't be called */
        return 1;
    }

    /* void __iomem *devm_ioremap(struct device *dev, resource_size_t offset, resource_size_t size) */
    info->phys_addr = (phys_addr_t)regs_get_kernel_argument(regs, 1);
    info->size = (size_t)regs_get_kernel_argument(regs, 2);

    return 0;
}
NOKPROBE_SYMBOL(devm_ioremap_entry);

static int devm_ioremap_return(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;

    info->virt_addr = (void __iomem *)regs_return_value(regs);

    register_emu_io(info->virt_addr, info->phys_addr, info->size);

    /* override with return emulation address */
    regs_set_return_value(regs, (unsigned long)virt2emu(info->virt_addr));
    instruction_pointer_set(regs, procedure_link_pointer(regs));

    return 0;
}
NOKPROBE_SYMBOL(devm_ioremap_return);

static struct kretprobe devm_ioremap_kretprobe = {
    .kp.symbol_name = "devm_ioremap",
	.entry_handler	= devm_ioremap_entry,
    .handler		= devm_ioremap_return,
	.data_size		= sizeof(struct iomap_info),
	/* Probe up to 10 instances concurrently. */
	.maxactive		= 10,
};

static int devm_ioremap_resource_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;
    unsigned long caller_address = procedure_link_pointer(regs);
    struct resource *res;

    if ((info->cnmod = lookup_module_by_text_address(caller_address)) == NULL) {
        /* The caller is not our modules, ignore */
        /* return 1, kret handler won't be called */
        return 1;
    }

    /* void __iomem *devm_ioremap_resource(struct device *dev, const struct resource *res) */
    res = (struct resource *)regs_get_kernel_argument(regs, 1);
    info->phys_addr = res->start;
    info->size = resource_size(res);

    return 0;
}
NOKPROBE_SYMBOL(devm_ioremap_resource_entry);

static int devm_ioremap_resource_return(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;

    info->virt_addr = (void __iomem *)regs_return_value(regs);

    register_emu_io(info->virt_addr, info->phys_addr, info->size);

    /* override with return emulation address */
    regs_set_return_value(regs, (unsigned long)virt2emu(info->virt_addr));
    instruction_pointer_set(regs, procedure_link_pointer(regs));

    return 0;
}
NOKPROBE_SYMBOL(devm_ioremap_resource_return);

static struct kretprobe devm_ioremap_resource_kretprobe = {
    .kp.symbol_name = "devm_ioremap_resource",
	.entry_handler	= devm_ioremap_resource_entry,
    .handler		= devm_ioremap_resource_return,
	.data_size		= sizeof(struct iomap_info),
	/* Probe up to 10 instances concurrently. */
	.maxactive		= 10,
};

static int iounmap_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;
    unsigned long caller_address = procedure_link_pointer(regs);

    if ((info->cnmod = lookup_module_by_text_address(caller_address)) == NULL) {
        /* The caller is not our modules, ignore */
        /* return 1, kret handler won't be called */
        return 1;
    }

    /*  void iounmap(volatile void __iomem *addr) */
    info->virt_addr = (void __iomem *)regs_get_kernel_argument(regs, 0);
    info->virt_addr = emu2virt(info->virt_addr);
    pt_regs_write_reg(regs, 0, (unsigned long)info->virt_addr);

    unregister_emu_io(info->virt_addr);

    return 0;
}
NOKPROBE_SYMBOL(iounmap_entry);

static struct kretprobe iounmap_kretprobe = {
    .kp.symbol_name = "iounmap",
	.entry_handler	= iounmap_entry,
	.data_size		= sizeof(struct iomap_info),
	/* Probe up to 10 instances concurrently. */
	.maxactive		= 10,
};

static int devm_iounmap_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct iomap_info *info = (struct iomap_info *)ri->data;
    unsigned long caller_address = procedure_link_pointer(regs);

    if ((info->cnmod = lookup_module_by_text_address(caller_address)) == NULL) {
        /* The caller is not our modules, ignore */
        /* return 1, kret handler won't be called */
        return 1;
    }

    /* void devm_iounmap(struct device *dev, void __iomem *addr) */
    info->virt_addr = (void __iomem *)regs_get_kernel_argument(regs, 1);
    info->virt_addr = emu2virt(info->virt_addr);
    pt_regs_write_reg(regs, 1, (unsigned long)info->virt_addr);

    unregister_emu_io(info->virt_addr);

    return 0;
}
NOKPROBE_SYMBOL(devm_iounmap_entry);

static struct kretprobe devm_iounmap_kretprobe = {
    .kp.symbol_name = "devm_iounmap",
	.entry_handler	= devm_iounmap_entry,
	.data_size		= sizeof(struct iomap_info),
	/* Probe up to 10 instances concurrently. */
	.maxactive		= 10,
};

struct exception_info {
    struct module_info *cnmod;

    struct pt_regs *regs;
};

static int fixup_exception_entry(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct exception_info *info = (struct exception_info *)ri->data;
    unsigned long pc;

    /*
     * see arch/arm64/mm/extable.c
     * int fixup_exception(struct pt_regs *regs)
     */
    info->regs = (struct pt_regs *)regs_get_kernel_argument(regs, 0);
    pc = (info->regs)->pc;

    if ((info->cnmod = lookup_module_by_text_address(pc)) == NULL) {
        /* The caller is not our modules, ignore */
        /* return 1, kret handler won't be called */
        return 1;
    }

    return 0;
}
NOKPROBE_SYMBOL(fixup_exception_entry);

static int fixup_exception_return(struct kretprobe_instance *ri, struct pt_regs *regs)
{
	struct exception_info *info = (struct exception_info *)ri->data;
    uint64_t current_pc, next_pc;
    uint32_t insn;
    mmio_op_t mmio;
    int status = STATUS_FAIL;

    current_pc = instruction_pointer(info->regs);
    insn = le32_to_cpu(*(uint32_t *)current_pc);

    mmio.regs = info->regs;
    if (decode_aarch64_insn(insn, &mmio) == STATUS_FAIL) {
        /* 不是ldr str指令导致的异常, 不处理 */
        pr_err("Invalid instruction decoding\n");
        return 0;
    }

    pr_debug("[addr]0x%llx [dir]=%s, size=%d, reg[%d]\n", (uint64_t)mmio.address, mmio.is_store?"write":"read", mmio.access_size, mmio.regx);

    if (!g_emulation_handle)
        return 0;

    mmio.address = emu2virt((void __iomem *)mmio.address);

    /* Call the emulation handler */
    status = g_emulation_handle(&mmio);
    if (status != STATUS_OK) {
        return 0;
    }

    /* For read operations, put read value back in stack frame */
    if (!mmio.is_store) {
        set_x_reg(mmio.regs, mmio.regx, mmio.value);
    }

    /*
     * 返回到异常指令下一条指令,如:
     * ldr  w1, [x2]    =====> 指令试图读取x2地址的上32bit数据到w1寄存器，发生data abort,进入异常处理
     * add  w1, w1, #1  <===== 在异常处理流程中, pc = pc + 4,即可返回到这一条指令继续执行
     */
    next_pc = current_pc + WORD_INSTRUCTION_LENGTH;
    instruction_pointer_set(info->regs, next_pc);

    /* override fixup_exception() with return 1 */
    regs_set_return_value(regs, 1);
    instruction_pointer_set(regs, procedure_link_pointer(regs));

    return 0;
}
NOKPROBE_SYMBOL(fixup_exception_return);

static struct kretprobe fixup_exception_kretprobe = {
    .kp.symbol_name = "fixup_exception",
	.entry_handler	= fixup_exception_entry,
    .handler		= fixup_exception_return,
	.data_size		= sizeof(struct exception_info),
	/* Probe up to 10 instances concurrently. */
	.maxactive		= 10,
};

#define NR_KRETPROBES   6

static struct kretprobe *g_kretprobes[NR_KRETPROBES] = {
    &ioremap_kretprobe,
    &iounmap_kretprobe,

    &devm_ioremap_kretprobe,
    &devm_ioremap_resource_kretprobe,
    &devm_iounmap_kretprobe,

    &fixup_exception_kretprobe,
};

int trap_io_init(emulation_handle_t *fn)
{
	int ret;

    ei_debugfs_init();

	ret = register_kretprobes(g_kretprobes, NR_KRETPROBES);
	if (ret < 0) {
		pr_err("register_kretprobe failed, returned %d\n", ret);
		return -1;
	}

    g_emulation_handle = fn;

	return 0;
}

void trap_io_deinit(void)
{
    int i;

	unregister_kretprobes(g_kretprobes, NR_KRETPROBES);

	/* nmissed > 0 suggests that maxactive was set too low. */
    for (i = 0; i < NR_KRETPROBES; i++)
        pr_debug("Missed probing %d instances of %s\n",
		    g_kretprobes[i]->nmissed, g_kretprobes[i]->kp.symbol_name);

    g_emulation_handle = NULL;
}

